package redis.inaction.ratelimiter.filter.ratelimit;

import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.cloud.gateway.filter.ratelimit.AbstractRateLimiter;
import org.springframework.cloud.gateway.filter.ratelimit.RedisRateLimiter;
import org.springframework.cloud.gateway.support.ConfigurationService;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.data.redis.core.ReactiveStringRedisTemplate;
import org.springframework.data.redis.core.script.RedisScript;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import java.util.*;
import java.util.concurrent.atomic.AtomicBoolean;

@Slf4j
public class ApiRedisRateLimiter extends AbstractRateLimiter<RedisRateLimiter.Config> implements ApplicationContextAware {

    private final AtomicBoolean initialized = new AtomicBoolean(false);

    @Autowired
    private RedisRateLimiter redisRateLimiter;

    private ReactiveStringRedisTemplate redisTemplate;

    private RedisScript<List<Long>> script;

    public ApiRedisRateLimiter(ReactiveStringRedisTemplate redisTemplate, RedisScript<List<Long>> script,
                               ConfigurationService configurationService) {
        super(RedisRateLimiter.Config.class, RedisRateLimiter.CONFIGURATION_PROPERTY_NAME, configurationService);
        this.redisTemplate = redisTemplate;
        this.script = script;
        this.initialized.compareAndSet(false, true);
    }

    static List<String> getKeys(String id) {
        String prefix = "request_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    public Mono<Response> isAllowed(int replenishRate, int burstCapacity, int requestedTokens, String apiCode){
        try {
            List<String> keys = getKeys(apiCode);
            List<String> scriptArgs = Arrays.asList(replenishRate + "", burstCapacity + "", "", requestedTokens + "");
            Flux<List<Long>> flux = this.redisTemplate.execute(this.script, keys, scriptArgs);
            return flux.onErrorResume(throwable -> {
                if (log.isDebugEnabled()) {
                    log.debug("Error calling rate limiter lua", throwable);
                }
                return Flux.just(Arrays.asList(1L, -1L));
            }).reduce(new ArrayList<Long>(), (longs, l) -> {
                longs.addAll(l);
                return longs;
            }).map(results -> {
                boolean allowed = results.get(0) == 1L;
                Long tokensLeft = results.get(1);
                Response response = new Response(allowed, getHeaders(replenishRate, burstCapacity, requestedTokens, tokensLeft));
                if (log.isDebugEnabled()) {
                    log.debug("response: " + response);
                }
                return response;
            });
        }
        catch (Exception e) {
            log.error("Error determining if user allowed from redis", e);
        }
        return Mono.just(new Response(true, getHeaders(replenishRate, burstCapacity, requestedTokens, -1L)));
    }

    @Override
    public Mono<Response> isAllowed(String routeId, String id) {
        return redisRateLimiter.isAllowed(routeId, id);
    }

    @Override
    public void setApplicationContext(ApplicationContext context) throws BeansException {
        if (initialized.compareAndSet(false, true)) {
            if (this.redisTemplate == null) {
                this.redisTemplate = context.getBean(ReactiveStringRedisTemplate.class);
            }
            this.script = context.getBean(RedisRateLimiter.REDIS_SCRIPT_NAME, RedisScript.class);
            if (context.getBeanNamesForType(ConfigurationService.class).length > 0) {
                setConfigurationService(context.getBean(ConfigurationService.class));
            }
        }
    }
    public Map<String, String> getHeaders(int replenishRate, int burstCapacity, int requestedTokens, Long tokensLeft) {
        Map<String, String> headers = new HashMap<>();
        if (redisRateLimiter.isIncludeHeaders()) {
            headers.put(redisRateLimiter.getRemainingHeader(), tokensLeft.toString());
            headers.put(redisRateLimiter.getReplenishRateHeader(), String.valueOf(replenishRate));
            headers.put(redisRateLimiter.getBurstCapacityHeader(), String.valueOf(burstCapacity));
            headers.put(redisRateLimiter.getRequestedTokensHeader(), String.valueOf(requestedTokens));
        }
        return headers;
    }

}
